import numpy as np
import mesa
import random
import seaborn as sns
def register_agent(agent_class):
  return agent_class
'''
Basic food producer
'''
@register_agent
class Producer(mesa.Agent):
  def __init__(self,unique_id,world,agent_class,energy=1,max_energy=10,grow_speed=0.01):
    super().__init__(world)
    self.unique_id = unique_id
    self.agent_class = agent_class
    self.energy = energy
    self.max_energy = max_energy
    self.grow_speed = grow_speed
    self.action = [0,0]
    self.speed = 0
  def step(self):
    if self.energy < self.max_energy:
      self.energy += self.grow_speed
  def arrest(self,amout):
    eat_amount = min(self.energy,amout)
    self.energy -= eat_amount
    return eat_amount
'''
Test agent, just for testing
'''
@register_agent
class Animal_Test(mesa.Agent):
  def __init__(self,unique_id,world,agent_class,prey=[],predator=[],cowoker=[],energy=1,eat_speed=2):
    super().__init__(world)
    self.unique_id = unique_id
    self.agent_class = agent_class
    self.prey = prey
    self.predator = predator
    self.cowoker = cowoker
    self.energy = energy
    self.eat_speed = eat_speed
    self.action = [0,0]
    self.speed = 0
  def step(self):
    self.random_move()
    cell_contents = self.model.grid.get_cell_list_contents([self.pos])
    for agent in cell_contents:
      if agent.agent_class in self.prey:
        eat_amount = agent.arrest(self.eat_speed)
  def random_move(self):
    if np.random.rand()<0.5:
      possible_moves = self.model.grid.get_neighborhood(self.pos,moore=True,include_center=False)
      new_pos = random.choice(possible_moves)
      self.model.grid.move_agent(self,new_pos)
    else:
      pass
  def arrest(self,amout):
    eat_amount = min(self.energy,amout)
    return eat_amount
'''
Basic animal model
'''
@register_agent
class Animal(mesa.Agent):
  def __init__(self,unique_id,world,brain,agent_class,prey=[],predator=[],cowoker=[],energy=1,max_energy=100000,eat_speed=2,observation_radius=6,base_energy_decay=0.01,move_energy_cost=0.05):
    super().__init__(world)
    self.brain = brain
    self.unique_id = unique_id
    self.agent_class = agent_class
    self.prey = prey
    self.predator = predator
    self.cowoker = cowoker
    self.energy = energy
    self.max_energy = max_energy
    self.eat_speed = eat_speed
    self.observation_radius = observation_radius
    self.observation_field = observation_radius*2+1
    self.base_energy_decay = base_energy_decay
    self.move_energy_cost = move_energy_cost
    self.reward = 0
    self.steps = 0
    self.observation = 0
    self.reward = 0
    self.action_idx = 0
    self.next_observation = 0
    self.action = [0,0]
    self.speed = 1
  def step(self):
    self.steps += 1
    observation = self.observe()
    self.memory()
    self.die()
    self.energy_decay()
    action = self.think(observation)
    self.move(action)
    self.eat()
    self.reproduction()
  def memory(self):
    if self.steps>1:
      self.brain.remember(self.observation,self.action_idx,self.reward,self.next_observation,self.energy<=0)
    self.reward = 0
  def energy_decay(self):
    self.energy -= self.base_energy_decay*self.energy 
    self.reward -= self.base_energy_decay*self.energy 
  def observe(self):
    observation = np.zeros((3,self.observation_field,self.observation_field))
    for neighbor in self.model.grid.get_neighbors(self.pos,moore=True,radius=self.observation_radius):
      relative_pos = (np.array(neighbor.pos)-np.array(self.pos)).astype(int)
      if neighbor.agent_class in self.prey:
        observation[0,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
      elif neighbor.agent_class in self.cowoker:
        observation[1,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
      elif neighbor.agent_class in self.predator:
        observation[2,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
    observation[1,self.observation_radius,self.observation_radius] = self.energy
    for i in range(self.observation_field):
      for j in range(self.observation_field):
        pos = np.array(self.pos)+np.array([i,j])-np.array([self.observation_radius,self.observation_radius])
        if self.model.grid.out_of_bounds(pos):
          observation[:,i,j] = -1
    #observation = (observation - self.model.observation_stats['mean'])/(self.model.observation_stats['std'] + 1e-8)
    self.observation = self.next_observation*1
    self.next_observation = observation*1
    return observation
  def think(self,observation):
    action_idx = self.brain.choose_action(observation)
    action = self.brain.action_map[action_idx]
    self.action_idx = action_idx
    return action
  def move(self,action):
    action = np.array(action)
    self.action = action
    if np.sum(abs(action))>0:
      self.energy -= self.move_energy_cost
      self.reward -= self.move_energy_cost
    new_pos = (
      self.pos[0] + action[0],
      self.pos[1] + action[1]
    )
    if self.model.grid.out_of_bounds(new_pos):
      new_pos = self.pos
    else:
      self.model.grid.move_agent(self, new_pos)
  def eat(self):
    cell_contents = self.model.grid.get_cell_list_contents([self.pos])
    for i in cell_contents:
      if i.agent_class in self.prey:
        eat_amount = i.arrest(self.eat_speed)
        self.energy += eat_amount
        self.reward += eat_amount
        self.energy = min(self.energy,self.max_energy)
  def reproduction(self):
    #if self.energy >= parameter_sheep_born_size*2:
    #    self.energy -= parameter_sheep_born_size
    #    self.model.add_sheep()
    return
  def arrest(self,amout):
    eat_amount = min(self.energy,amout)
    self.energy -= eat_amount
    self.reward -= eat_amount
    return eat_amount
  def die(self):
    if self.energy <= 0:
      self.model.remove_agent(self)
'''
Basic animal model
'''
@register_agent
class Animal(mesa.Agent):
  def __init__(self,unique_id,world,brain,agent_class,prey=[],predator=[],cowoker=[],energy=1,max_energy=100000,eat_speed=2,observation_radius=6,base_energy_decay=0.01,move_energy_cost=0.05):
    super().__init__(world)
    self.brain = brain
    self.unique_id = unique_id
    self.agent_class = agent_class
    self.prey = prey
    self.predator = predator
    self.cowoker = cowoker
    self.energy = energy
    self.max_energy = max_energy
    self.eat_speed = eat_speed
    self.observation_radius = observation_radius
    self.observation_field = observation_radius*2+1
    self.base_energy_decay = base_energy_decay
    self.move_energy_cost = move_energy_cost
    self.reward = 0
    self.steps = 0
    self.observation = 0
    self.reward = 0
    self.action_idx = 0
    self.next_observation = 0
    self.action = [0,0]
    self.speed = 1
  def step(self):
    self.steps += 1
    observation = self.observe()
    self.memory()
    self.die()
    self.energy_decay()
    action = self.think(observation)
    self.move(action)
    self.eat()
    self.reproduction()
  def memory(self):
    if self.steps>1:
      self.brain.remember(self.observation,self.action,self.reward,self.next_observation,self.energy<=0)
    self.reward = 0
  def energy_decay(self):
    self.energy -= self.base_energy_decay*self.energy 
    self.reward -= self.base_energy_decay*self.energy 
  def observe(self):
    observation = np.zeros((3,self.observation_field,self.observation_field))
    for neighbor in self.model.grid.get_neighbors(self.pos,moore=True,radius=self.observation_radius):
      relative_pos = (np.array(neighbor.pos)-np.array(self.pos)).astype(int)
      if neighbor.agent_class in self.prey:
        observation[0,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
      elif neighbor.agent_class in self.cowoker:
        observation[1,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
      elif neighbor.agent_class in self.predator:
        observation[2,self.observation_radius+relative_pos[0],self.observation_radius+relative_pos[1]] = neighbor.energy
    observation[1,self.observation_radius,self.observation_radius] = self.energy
    for i in range(self.observation_field):
      for j in range(self.observation_field):
        pos = np.array(self.pos)+np.array([i,j])-np.array([self.observation_radius,self.observation_radius])
        if self.model.grid.out_of_bounds(pos):
          observation[:,i,j] = -1
    #observation = (observation - self.model.observation_stats['mean'])/(self.model.observation_stats['std'] + 1e-8)
    self.observation = self.next_observation*1
    self.next_observation = observation*1
    return observation
  def think(self,observation):
    return self.brain.choose_action(observation)
  def move(self,action):
    action = np.array(action)
    self.action = action
    action = np.round(action[:2]/np.linalg.norm(action[:2])*action[2],0).astype(int)
    length = np.linalg.norm(action)
    if np.sum(abs(action))>0:
      self.energy -= self.move_energy_cost*length
      self.reward -= self.move_energy_cost*length
    new_pos = (
      self.pos[0] + action[0],
      self.pos[1] + action[1]
    )
    if self.model.grid.out_of_bounds(new_pos):
      new_pos = self.pos
    else:
      self.model.grid.move_agent(self, new_pos)
  def eat(self):
    cell_contents = self.model.grid.get_cell_list_contents([self.pos])
    for i in cell_contents:
      if i.agent_class in self.prey:
        eat_amount = i.arrest(self.eat_speed)
        self.energy += eat_amount
        self.reward += eat_amount
        self.energy = min(self.energy,self.max_energy)
  def reproduction(self):
    #if self.energy >= parameter_sheep_born_size*2:
    #    self.energy -= parameter_sheep_born_size
    #    self.model.add_sheep()
    return
  def arrest(self,amout):
    eat_amount = min(self.energy,amout)
    self.energy -= eat_amount
    self.reward -= eat_amount
    return eat_amount
  def die(self):
    if self.energy <= 0:
      self.model.remove_agent(self)